#!/usr/bin/env python
# -*-coding: utf-8-*-
"""
@Author:      chen ming
@Date:        2025/5/1 
@Time:        10:14
@Description: 
"""

import numpy as np
from scipy.integrate import solve_ivp


class ROVDynamics:
    def __init__(self, mass=100, volume=0.1, thruster_positions=None, thruster_directions=None):
        """ROV六自由度动力学模型
        Args:
            mass (float): 质量 (kg)
            volume (float): 排水体积 (m³)
            thruster_positions (np.array): 8×3推进器安装位置 (m)
            thruster_directions (np.array): 8×3推进器方向单位向量
        """
        # 基本参数
        self.mass = mass
        self.volume = volume
        self.gravity = 9.81
        self.water_density = 1025  # 海水密度 kg/m³

        # 惯性矩阵构建
        self.I_rot = np.diag([10, 10, 5])       # 旋转惯量 (3x3)
        self.I_trans = np.eye(3) * mass         # 平移惯量 (3x3)
        self.M = np.block([[self.I_trans, np.zeros((3, 3))],
                           [np.zeros((3, 3)), self.I_rot]])  # 6x6惯性矩阵

        # 推进器分配矩阵
        self.B = self.calculate_thruster_matrix(thruster_positions, thruster_directions)

        # 初始化状态变量 [x,y,z,phi,theta,psi, u,v,w,p,q,r]
        self.state = np.zeros(12)

    def calculate_thruster_matrix(self, positions, directions):
        """计算推进器分配矩阵 (6x8)"""
        if positions is None or directions is None:
            return np.zeros((6, 8))  # 默认空矩阵

        B = np.zeros((6, 8))
        for i in range(8):
            B[0:3, i] = directions[i]
            B[3:6, i] = np.cross(positions[i], directions[i])
        return B

    def coriolis_matrix(self, nu):
        """计算科里奥利矩阵 (6x6)"""
        omega = nu[3:]  # 角速度分量 [p,q,r]
        S = self.skew_symmetric(omega)

        # 分块构建科里奥利矩阵
        C_rot = -S @ self.I_rot  # 核心旋转分量
        return np.block([[np.zeros((3, 3)), np.zeros((3, 3))],
                         [np.zeros((3, 3)), C_rot]])

    @staticmethod
    def skew_symmetric(v):
        """生成三维反对称矩阵"""
        return np.array([[0, -v[2], v[1]],
                         [v[2], 0, -v[0]],
                         [-v[1], v[0], 0]])

    @staticmethod
    def damping_matrix(nu):
        """简化的线性阻尼矩阵 (6x6)"""
        D_lin = np.diag([100, 100, 200, 50, 50, 30])  # 线性阻尼系数
        return D_lin

    def restoring_forces(self):
        """计算静力恢复力(重力+浮力)"""
        weight = self.mass * self.gravity
        buoyancy = self.water_density * self.volume * self.gravity

        # 假设重心与浮心在z轴有偏移
        z_g = 0.1  # 重心垂向位置
        z_b = 0.2  # 浮心垂向位置

        return np.array([0, 0, buoyancy - weight,
                         0, 0, (z_g - z_b) * (buoyancy - weight)])

    def dynamics_model(self, t, state, tau):
        """动力学微分方程定义"""
        eta = state[0:6]  # 位置/姿态
        nu = state[6:12]  # 速度/角速度

        # 矩阵计算
        C = self.coriolis_matrix(nu)
        D = self.damping_matrix(nu)
        g = self.restoring_forces()

        # 系统加速度
        nu_dot = np.linalg.inv(self.M) @ (
                self.B @ tau - (C + D) @ nu - g
        )

        # 运动学方程（简化处理姿态）
        eta_dot = nu  # 注意：实际应使用旋转矩阵进行转换

        return np.concatenate((eta_dot, nu_dot))


# ===================== 测试用例 ========================#
if __name__ == "__main__":
    # 推进器配置示例（前四个推进器）
    positions = np.array([[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0],
                          [0, 0, 1], [0, 0, -1], [0, 0, 1], [0, 0, -1]])  # 8x3
    directions = np.array([[1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 1, 0],
                           [0, 0, 1], [0, 0, 1], [0, 0, 1], [0, 0, 1]])  # 8x3

    # 初始化模型
    rov = ROVDynamics(mass=150,
                      thruster_positions=positions,
                      thruster_directions=directions)

    # 验证矩阵维度
    print("惯性矩阵维度:", rov.M.shape)  # 应显示(6,6)
    C_test = rov.coriolis_matrix(np.array([0, 0, 0, 1, 2, 3]))
    print("科里奥利矩阵维度:", C_test.shape)  # 应显示(6,6)


    # 模拟测试
    def simulator(t, y):
        return rov.dynamics_model(t, y, tau=np.zeros(8))  # 无推力输入


    sol = solve_ivp(simulator, [0, 10],
                    np.zeros(12),
                    method='RK45')

    # 绘制结果
    import matplotlib.pyplot as plt

    plt.figure(figsize=(10, 6))
    plt.plot(sol.t, sol.y[0], label='x position')
    plt.plot(sol.t, sol.y[1], label='y position')
    plt.xlabel('Time (s)')
    plt.ylabel('Position (m)')
    plt.legend()
    plt.grid(True)
    plt.show()
